#include "tcp_connection.hh"

#include <iostream>

// Dummy implementation of a TCP connection

// For Lab 4, please replace with a real implementation that passes the
// automated checks run by `make check`.

using namespace std;

size_t TCPConnection::remaining_outbound_capacity() const {
    //cerr << "\n-- " << _linger_after_streams_finish << ": remaining capacity is " << _sender.stream_in().remaining_capacity() << endl;
    return _sender.stream_in().remaining_capacity();
}

size_t TCPConnection::bytes_in_flight() const {
    return _sender.bytes_in_flight();
}

size_t TCPConnection::unassembled_bytes() const {
    return _receiver.unassembled_bytes();
}

size_t TCPConnection::time_since_last_segment_received() const {
    return _cur_time - _last_segment_recved_time;
}

void TCPConnection::segment_received(const TCPSegment &seg) {
    _last_segment_recved_time = _cur_time;
    _update_state();
    const TCPHeader &header = seg.header();
    if(header.rst){
//        cerr << "\n-- " << (_linger_after_streams_finish? "主动方" : "被动方") << " recv RST," << header.seqno.raw_value() << endl;
        if(_cur_state != TCPState::State::SYN_RCVD)
            _rst_handler();
        else
            _syn_recvd_rst_handler();
        return;
    }

//    cerr << "\n-- " << _linger_after_streams_finish << ": segment_received, state is " << SC_UINT(_cur_state) << "\n";

    _newest_is_pure_ack_seg = header.ack && seg.length_in_sequence_space() == 0;
    uint32_t old_recv_ackno;
    bool has_value = false;
    const auto &&recv_ackno = _receiver.ackno();
    if(recv_ackno.has_value()) {
        has_value = true;
        old_recv_ackno = recv_ackno.value().raw_value();
    }

    if(!_fin_record.fin_recvd && header.fin) {
        _fin_record.fin_recvd = true;
        _fin_record.fin_abs_index = unwrap(header.seqno + static_cast<uint32_t>(seg.length_in_sequence_space() - 1),
                                          _isn, _receiver.stream_out().bytes_written()) - 1;
    }

//    if(!_fin_recvd && header.fin) {
//        _fin_recvd = true;
//    }

    if(not _newest_is_pure_ack_seg)
        _receiver.segment_received(seg);

//    cerr << "\n--" << (_linger_after_streams_finish? "主动方" : "被动方") << " recv seg: "
//         << "对方的win: " << _receiver.window_size() << ", state is " << SC_UINT(_cur_state) << endl;

    if(header.ack) {
        _sender.ack_received(WrappingInt32(header.ackno), header.win);
    }


    if(_linger_after_streams_finish){
//        cerr << "\n--" << "主动方：old_recv_ackno is " << old_recv_ackno << ", recv seqno is " << header.seqno.raw_value() << endl;
    }
    if(has_value && not _newest_is_pure_ack_seg && _receiver.ackno().value().raw_value() == old_recv_ackno){
//        cerr << "\n--" << (_linger_after_streams_finish? "主动方" : "被动方") << " recv a out date seg: " << endl;
        _sender.send_empty_segment();
        _send_helper();
        return;
    }

    RecvOp recv{};
    recv.recvs = (header.syn ? RecvMask::RECV_SYN : 0) | (header.fin ? RecvMask::RECV_FIN : 0) |
                 (header.ack ? RecvMask::RECV_ACK : 0) | (header.rst ? RecvMask::RECV_RST : 0);
    InfoCodeWrap code_wrap(_cur_state, recv);
    if (_code_handler_map.find(code_wrap.uint_value) != _code_handler_map.end()) {
        auto &&handler = bind(_code_handler_map.at(code_wrap.uint_value), this);
        handler();
        _send_helper();
    }

    //_send_helper();
}

bool TCPConnection::active() const {
    auto &&sender_summary = TCPState::state_summary(_sender);
    auto &&recv_summary = TCPState::state_summary(_receiver);
    return !(sender_summary == TCPSenderStateSummary::ERROR && recv_summary == TCPReceiverStateSummary::ERROR)
        && !(!_linger_after_streams_finish && sender_summary == TCPSenderStateSummary::FIN_ACKED
                && recv_summary == TCPReceiverStateSummary::FIN_RECV);
}

size_t TCPConnection::write(const string &data) {
    if(_sender.stream_in().input_ended())
        return 0;
    size_t &&count = _sender.stream_in().write(data);
//    cerr << "\n-- " << _linger_after_streams_finish << ": write " << count << endl;
    if(count) {
        _sender.fill_window();
        _send_helper();
    }
    return count;
}

//! \param[in] ms_since_last_tick number of milliseconds since the last call to this method
void TCPConnection::tick(const size_t ms_since_last_tick) {
    _cur_time += ms_since_last_tick;
    _sender.tick(ms_since_last_tick);
    if(_linger_after_streams_finish && _timewait_start < _cur_time && _cur_time - _timewait_start >= _Linger_Time){
        // TimeWait -> Closed
        _linger_after_streams_finish = false;
        _timewait_start = SIZE_MAX;
        _update_state();
    }else if(_sender.consecutive_retransmissions() > TCPConfig::MAX_RETX_ATTEMPTS){
        _update_state();
        // 重传次数超限，需要关闭链接，并发送RST
        _sender.send_empty_segment();
        auto &&seg = move(_sender.segments_out().back());
        seg.header().rst = true;
        while(not _segments_out.empty())
            _segments_out.pop();
        _segments_out.emplace(seg);

//        cerr << "\n-- " << (_linger_after_streams_finish? "主动方" : "被动方") << "state is " << SC_UINT(_cur_state)
//             << ", retran beyond max, then send RST\n";

        _rst_handler();
    }else if(not _sender.segments_out().empty()){
        _send_helper();
    }
//    else if(_cur_state == TCPState::State::CLOSE_WAIT){
//        // 进入last_ack阶段
//        if(not _sender.stream_in().buffer_empty())
//            _sender.fill_window();
//
//        end_input_stream();
//    }
}

void TCPConnection::end_input_stream() {
//    cerr << "\n-- " << _linger_after_streams_finish << ": end_input_stream\n";
    _sender.stream_in().end_input();
    _sender.fill_window();
    _send_helper();
}

void TCPConnection::connect() {
    //_update_state();
    //printf("connect\n");
    RecvOp op{};
    op.ops = _cur_state == TCPState::State::CLOSED? OperationMask::OP_CONNECT : OperationMask::OP_SEND;
    InfoCodeWrap code_wrap(_cur_state, op);
    //printf("state is %u, op is %u, uintcode is %u\n", code_wrap.info_code.state_info, code_wrap.info_code.recv_op_info, code_wrap.uint_value);
    if(_code_handler_map.find(code_wrap.uint_value) != _code_handler_map.end()) {
        auto &&handler = bind(_code_handler_map.at(code_wrap.uint_value), this);
        handler();
    }
}

TCPConnection::~TCPConnection() {
    try {
        if (active()) {
            cerr << "Warning: Unclean shutdown of TCPConnection\n";

            // Your code here: need to send a RST segment to the peer
            _sender.send_empty_segment();
            auto &seg = _sender.segments_out().back();
            seg.header().rst = true;
            _send_helper();
//            cerr << "\n-- " << (_linger_after_streams_finish? "主动方" : "被动方") << "send RST\n";
            _rst_handler();
        }
    } catch (const exception &e) {
        std::cerr << "Exception destructing TCP FSM: " << e.what() << std::endl;
    }
}

/** Handler **/
inline void TCPConnection::_closed_connect_handler() {
    if(_sender.next_seqno_absolute())
        return;
    _sender.fill_window();
    auto &outbound = _sender.segments_out();
    auto &seg = outbound.back();
    auto &header = seg.header();
    header.win = _receiver.window_size();
    _isn = header.seqno;
    _segments_out.emplace(move(seg));
    outbound.pop();
//    cerr << "\n-- _closed_connect_handler: header is " << _segments_out.front().header().to_string() << endl;
}

inline void TCPConnection::_listen_send_handler() {
    //printf("listen_send_syn\n");
    // 发送syn
    _closed_connect_handler();
}

inline void TCPConnection::_listen_handshake2_handler() {
    //printf("listen_handshake2\n");
    // 发送syn+ack
    if(_sender.next_seqno_absolute())
        return;
    _sender.fill_window();
//    auto &outbound = _sender.segments_out();
//    auto &seg = outbound.back();
//    auto &header = seg.header();
//    header.win = _receiver.window_size();
//    header.ack = true;
//    header.ackno = _receiver.ackno().value();
//    while(not outbound.empty()){
//        _segments_out.emplace(move(outbound.front()));
//        outbound.pop();
//    }
}

inline void TCPConnection::_listen_close_handler() {
    // receiver.ackno().has_value()
//    cerr << "\n-- _listen_close_handler\n";
    _receiver.stream_out().end_input();

}

inline void TCPConnection::_syn_sent_sim_open_handler() {
    //printf("syn_sent_sim_open\n");
    // 发送syn+ack
    _sender.send_empty_segment();
    auto &seg = _sender.segments_out().back();
    auto &header = seg.header();
    header.syn = true;
    header.seqno = _isn;
}

inline void TCPConnection::_syn_sent_handshake3_handler() {
    //printf("syn_sent_handshake3\n");
//    cerr << "\n-- _syn_sent_handshake3_handler\n";
    // 发送ack
    if(not _sender.stream_in().buffer_empty())
        _sender.fill_window();
    else
        _sender.send_empty_segment();
}

inline void TCPConnection::_syn_sent_close_handler() {

}

inline void TCPConnection::_syn_recvd_ack_handler() {
    //printf("syn_recvd_ack\n");
    // 发送ack
    if(not _sender.stream_in().buffer_empty())
        _sender.fill_window();
    else if(not _newest_is_pure_ack_seg)
        _sender.send_empty_segment();
}

inline void TCPConnection::_syn_recvd_fin_ack_handler() {
    // 对方在第三次握手时，同时发送了fin和ack
//    cerr << "\n-- _syn_recvd_fin_ack_handler\n";
    _established_passive_close_handler();
}

inline void TCPConnection::_syn_recvd_rst_handler() {
    _receiver = TCPReceiver{_cfg.recv_capacity};
    _sender = TCPSender{_cfg.send_capacity, _cfg.rt_timeout, _cfg.fixed_isn};
    _segments_out = queue<TCPSegment>();
    _linger_after_streams_finish = true;
    _timewait_start = SIZE_MAX;
    _last_segment_recved_time = 0;
    _cur_time = 0;
    _update_state();
}

inline void TCPConnection::_syn_recvd_close_handler() {
    // 发fin
//    cerr << "\n-- _syn_recvd_close_handler\n";
    _sender.stream_in().end_input();
    _sender.fill_window();
}

inline void TCPConnection::_established_active_close_handler() {
    // 发fin
//    cerr << "\n-- _established_active_close_handler\n";
    _sender.stream_in().end_input();
    _sender.fill_window();
}

inline void TCPConnection::_established_passive_close_handler() {
    // 发ack
//    cerr << "\n-- _established_passive_close_handler\n";
    // 进入close_wait
    _receiver.stream_out().end_input();
    if(not _sender.stream_in().buffer_empty() || _sender.stream_in().eof())
        _sender.fill_window();
    else
        _sender.send_empty_segment();
    _linger_after_streams_finish = false;
    _update_state();
}

inline void TCPConnection::_established_ack_transmit_handler() {
    // 先填充outbound，再将seg一个个取出，加入ackno，win，最后放入_segments_out
    //printf("established_ack_transmit, %zu chars could be sent\n", _sender.stream_in().buffer_size());
    auto &sender_stream = _sender.stream_in();
    if(not sender_stream.buffer_empty() || sender_stream.eof())
        _sender.fill_window();
    else if(not _newest_is_pure_ack_seg)
        _sender.send_empty_segment();
}

inline void TCPConnection::_fin_wait_1_ack_handler() {
    if(not _newest_is_pure_ack_seg && TCPState::state_summary(_sender) != TCPSenderStateSummary::FIN_ACKED){
        _sender.send_empty_segment();
    }
}

// remove
inline void TCPConnection::_fin_wait_1_fin_handler() {
    _sender.send_empty_segment();
}

inline void TCPConnection::_fin_wait_1_fin_ack_handler() {
//    cerr << "\n-- _fin_wait_1_fin_ack_handler\n";
    _sender.send_empty_segment();
    // 收到的ack已对发送的fin确认
    if(TCPState::state_summary(_sender) == TCPSenderStateSummary::FIN_ACKED) {
        // 进入TimeWait
        _update_state();
        _timewait_start = _cur_time;
    }
}

inline void TCPConnection::_fin_wait_2_fin_handler() {
    // 进入TimeWait
//    cerr << "\n-- _fin_wait_2_fin_handler\n";
//    _receiver.stream_out().end_input();
//    _sender.send_empty_segment();
//    _update_state();
//    _timewait_start = _cur_time;
    //if(_receiver.stream_out().bytes_written() == _fin_record.fin_abs_index) {
    if(_receiver.is_finished()){
        _receiver.stream_out().end_input();
        _sender.send_empty_segment();
        _update_state();
        _timewait_start = _cur_time;
    }
}

inline void TCPConnection::_fin_wait_2_ack_handler() {
    // 在finwait2状态下收到一般数据
//    cerr << "\n-- _fin_wait_2_ack_handler\n";
    //if(_fin_record.fin_recvd && _receiver.stream_out().bytes_written() == _fin_record.fin_abs_index) {
    if(_fin_record.fin_recvd && _receiver.is_finished()){
        // _receiver接收完了所有数据
        _receiver.stream_out().end_input();
        _sender.send_empty_segment();
        _update_state();
        _timewait_start = _cur_time;
    }
//    if(_fin_recvd && _cur_state != TCPState::State::TIME_WAIT){
//        _receiver.stream_out().end_input();
//        _sender.send_empty_segment();
//        _update_state();
//        _timewait_start = _cur_time;
//    }
    else if(not _newest_is_pure_ack_seg)
        _sender.send_empty_segment();
}

inline void TCPConnection::_closing_ack_handler() {
    // 进入TimeWait
//    cerr << "\n-- _closing_ack_handler\n";
    _update_state();
    _timewait_start = _cur_time;
}

void TCPConnection::_time_wait_handler() {}

inline void TCPConnection::_close_wait_close_handler() {
//    cerr << "\n-- _close_wait_close_handler\n";
    _sender.stream_in().end_input();
    _sender.send_empty_segment();
}

inline void TCPConnection::_close_wait_ack_handler() {
//    cerr << "\n-- _close_wait_ack_handler: " << not _sender.stream_in().eof() << ", "
//         << (_sender.next_seqno_absolute() < _sender.stream_in().bytes_written() + 2) << endl;

    auto &sender_stream = _sender.stream_in();
    if(not sender_stream.buffer_empty() || sender_stream.eof())
        // sender的buf不为空要发，否则，sender的stream停止输入了，要发，此时肯定没有发过fin
        _sender.fill_window();
    else if(not _newest_is_pure_ack_seg)
        _sender.send_empty_segment();
//    auto &&segoutSize = _sender.segments_out().size();
//    _sender.fill_window();
//    if(not _newest_is_pure_ack_seg && _sender.segments_out().size() == segoutSize)
//        _sender.send_empty_segment();
}

inline void TCPConnection::_last_ack_handler() {
    _update_state();
//    cerr << "\n-- _last_ack_handler: now state is " << SC_UINT(_cur_state) << endl;
}

inline void TCPConnection::_send_helper() {
    auto &&ack_wrap = _receiver.ackno();
    auto &outbound = _sender.segments_out();
    size_t &&win_size = _receiver.window_size();
    if(ack_wrap.has_value()){
        auto &ackno = ack_wrap.value();
        while(not outbound.empty()){
            auto &seg = outbound.front();
            auto &header = seg.header();
//            if(seg.length_in_sequence_space() == 0){
//                cerr << "\n-- send_helper: empty seg's seqno is " << header.seqno << endl;
//            }
            header.ack = true;
            header.ackno = ackno;
            header.win = win_size;
            _segments_out.emplace(move(seg));
            outbound.pop();
        }
    }else{
        // 发送syn
        while(not outbound.empty()) {
//            cerr << "\n-- send_helper: 发送syn\n";
            auto &seg = outbound.front();
            seg.header().win = win_size;
            _segments_out.emplace(move(seg));
            outbound.pop();
        }
    }
}

inline void TCPConnection::_rst_handler() {
    _sender.stream_in().set_error();
    _receiver.stream_out().set_error();
    _linger_after_streams_finish = false;
}


//const std::vector<TCPState> TCPConnection::_official_state = []() -> std::vector<TCPState> {
//  std::vector<TCPState> v;
//  v.reserve(16);
//  for(int i = 0; i < 12; ++i){
//      v.emplace_back(static_cast<TCPState::State>(i));
//  }
//
//  return v;
//}();

const std::unordered_map<uint32_t, void (TCPConnection::*)()> TCPConnection::_code_handler_map
    = []() -> std::unordered_map<uint32_t, void (TCPConnection::*)()> {
      std::unordered_map<uint32_t, void (TCPConnection::*)()> umap;
      RecvOp ro{};
      ro.ops = OperationMask::OP_CONNECT;
      umap[InfoCodeWrap(TCPState::State::CLOSED, ro).uint_value] = &TCPConnection::_closed_connect_handler;

      ro.ops = OperationMask::OP_SEND;
      umap[InfoCodeWrap(TCPState::State::LISTEN, ro).uint_value] = &TCPConnection::_listen_send_handler;

      ro.recvs = RecvMask::RECV_SYN;
      umap[InfoCodeWrap(TCPState::State::LISTEN, ro).uint_value] = &TCPConnection::_listen_handshake2_handler;
      umap[InfoCodeWrap(TCPState::State::SYN_SENT, ro).uint_value] = &TCPConnection::_syn_sent_sim_open_handler;

      ro.recvs |= RecvMask::RECV_ACK;
      umap[InfoCodeWrap(TCPState::State::SYN_SENT, ro).uint_value] = &TCPConnection::_syn_sent_handshake3_handler;

      ro.recvs = RecvMask::RECV_ACK;
      umap[InfoCodeWrap(TCPState::State::SYN_RCVD, ro).uint_value] = &TCPConnection::_syn_recvd_ack_handler;
      umap[InfoCodeWrap(TCPState::State::FIN_WAIT_1, ro).uint_value] = &TCPConnection::_fin_wait_1_ack_handler;
      umap[InfoCodeWrap(TCPState::State::FIN_WAIT_2, ro).uint_value] = &TCPConnection::_fin_wait_2_ack_handler;
      umap[InfoCodeWrap(TCPState::State::CLOSING, ro).uint_value] = &TCPConnection::_closing_ack_handler;
      umap[InfoCodeWrap(TCPState::State::CLOSE_WAIT, ro).uint_value] = &TCPConnection::_close_wait_ack_handler;
      umap[InfoCodeWrap(TCPState::State::LAST_ACK, ro).uint_value] = &TCPConnection::_last_ack_handler;
      umap[InfoCodeWrap(TCPState::State::ESTABLISHED, ro).uint_value] = &TCPConnection::_established_ack_transmit_handler;

      ro.recvs |= RecvMask::RECV_FIN;
      umap[InfoCodeWrap(TCPState::State::ESTABLISHED, ro).uint_value] = &TCPConnection::_established_passive_close_handler;
      //umap[InfoCodeWrap(TCPState::State::FIN_WAIT_1, ro).uint_value] = &TCPConnection::_fin_wait_1_fin_handler;
      umap[InfoCodeWrap(TCPState::State::FIN_WAIT_2, ro).uint_value] = &TCPConnection::_fin_wait_2_fin_handler;
      umap[InfoCodeWrap(TCPState::State::FIN_WAIT_1, ro).uint_value] = &TCPConnection::_fin_wait_1_fin_ack_handler;
      umap[InfoCodeWrap(TCPState::State::SYN_RCVD, ro).uint_value] = &TCPConnection::_syn_recvd_fin_ack_handler;

      ro.ops = OperationMask::OP_CLOSE;
      umap[InfoCodeWrap(TCPState::State::LISTEN, ro).uint_value] = &TCPConnection::_listen_close_handler;
      umap[InfoCodeWrap(TCPState::State::SYN_SENT, ro).uint_value] = &TCPConnection::_syn_sent_close_handler;
      umap[InfoCodeWrap(TCPState::State::SYN_RCVD, ro).uint_value] = &TCPConnection::_syn_recvd_close_handler;
      umap[InfoCodeWrap(TCPState::State::ESTABLISHED, ro).uint_value] = &TCPConnection::_established_active_close_handler;
      umap[InfoCodeWrap(TCPState::State::CLOSE_WAIT, ro).uint_value] = &TCPConnection::_close_wait_close_handler;

      ro.recvs = RecvMask::RECV_RST;
      umap[InfoCodeWrap(TCPState::State::SYN_RCVD, ro).uint_value] = &TCPConnection::_syn_recvd_rst_handler;

      return umap;
    }();

int TCPConnection::num = 0;

